#include "numuCCPXSelection.hxx"
#include "numuCCSelection.hxx"
#include "baseSelection.hxx"
#include "CutUtils.hxx"
#include "DataClasses.hxx"

void numuCCPXSelection::DefineSteps()
{
    // Cuts must be added in the right order
    // last "true" means the step sequence is broken if cut is not passed (default is "false")
    // Add inclusive CC steps to all branches

  /*
   //for comparison numuCCSelection; only differ by ExternalFGD1lastlayersCut

    AddStep(StepBase::kCut,    "event quality",      new EventQualityCut(),           true);
    AddStep(StepBase::kCut,    "> 0 tracks ",        new TotalMultiplicityCut(),      true);
    AddStep(StepBase::kAction, "find leading tracks",new FindLeadingTracksAction());

    AddStep(StepBase::kAction, "find vertex",        new FindVertexAction());
    AddStep(StepBase::kAction, "fill_summary",       new FillSummaryAction_numuCC());
    AddStep(StepBase::kCut,    "quality+fiducial",   new TrackQualityFiducialCut(),   true);
    AddStep(StepBase::kAction, "find veto track",    new FindVetoTrackAction());
    AddStep(StepBase::kCut,    "veto",               new ExternalVetoCut());
    AddStep(StepBase::kAction, "find oofv track",    new FindOOFVTrackAction());
    AddStep(StepBase::kCut,    "External FGD1",      new ExternalFGD1lastlayersCut());
    AddStep(StepBase::kCut,    "muon PID",           new MuonPIDCut());

    SetBranchAlias(0,"trunk");
  */

    AddStep(StepBase::kCut,    "event quality",        new EventQualityCut(),           true);
    AddStep(StepBase::kCut,    "> 0 tracks ",          new TotalMultiplicityCut(),      true);  
    AddStep(StepBase::kAction, "find leading tracks",  new FindLeadingTracksAction());  

    // First branch is for cases 0,1,2
    AddStep(StepBase::kAction, "find vertex",        new FindVertexAction());  
    AddStep(StepBase::kCut,    "quality+fiducial",   new TrackQualityFiducialCut(),   true);  
    AddStep(StepBase::kAction, "find veto track",    new FindVetoTrackAction());
    AddStep(StepBase::kCut,    "veto",               new ExternalVetoCut(), kTRUE);
    AddStep(StepBase::kAction, "find oofv track",    new FindOOFVTrackAction());
    AddStep(StepBase::kCut,    "muon PID",           new MuonPIDCut(), kTRUE);

    //above same as numuCCSelection; only differ by ExternalFGD1lastlayersCut <<--

    //->
    AddStep(StepBase::kAction, "find michel e-",     new FindMichelElectronsAction()); //local
    AddStep(StepBase::kCut,    "no michel ele",      new MichelElectronCut(), kTRUE); //local

    //have to be before TwoThreeTrackerTracksCut since trackertracks size changes
    AddStep(StepBase::kAction, "find tracker tracks",  new FindTrackerTracksAction()); //local action; on std_numuccPX_actions::FindTracksWithTPC, std_numuccPX_actions::FindTracksWithFGDAndNoTPC

    AddStep(StepBase::kCut,    "track charge",  new TrackChargeCut(), kTRUE); //local

    //ok to move it here
    AddStep(StepBase::kCut, "find candidates",      new FindCandidatesCut(), kTRUE); //local action;

    AddStep(StepBase::kCut, "common vertex",      new CommonVertexCut(), kTRUE); //local

    //very important! without it no events shown in Scan
    SetBranchAlias(0,"muTPC-pTPC");
}

bool TrackChargeCut::Apply(const AnaEventB& event, AnaBoxB& box) const
{
    (void)event;

    AnaBoxCCPX * ccPXbox = static_cast<AnaBoxCCPX*>(&box);

    //only 2 or 3 tracks
    const UInt_t ntracker = ccPXbox->TrackerTracks.size();
    if(ntracker < 2 || ntracker > 3)
      return kFALSE;

    //require no missing charge measurement
    const UInt_t ntpc = ccPXbox->nPositiveTPCtracks+ccPXbox->nNegativeTPCtracks;
    if(ntpc != ntracker){
      return kFALSE;
    }

    // Exactly one positive track with good quality and in FV
    //ntot = 3, so either (1,2) or (2,1)
    if (ccPXbox->nPositiveTPCtracks == 1 || ccPXbox->nNegativeTPCtracks==1) 
        return true;
    else 
        return false;
}

bool FindCandidatesCut::Apply(const AnaEventB& event, AnaBoxB& box) const
{
    (void)event;

    AnaBoxCCPX * ccPXbox = static_cast<AnaBoxCCPX*>(&box); 

    const UInt_t ntrk = ccPXbox->TrackerTracks.size();
    if(ntrk>3){
      printf("FindCandidatesCut::Apply bad ntrk>3 %d\n", ntrk); exit(1);
    }

    Int_t charges[3];
    Int_t pids[3];
    Bool_t kused[3]={0,0,0};
    for(UInt_t itrk = 0; itrk < ntrk; itrk++){
      AnaTrackB* track = ccPXbox->TrackerTracks[itrk];
      charges[itrk]=(track->Charge>0);//pos: 1; neg: 0
      pids[itrk]=std_numuccPX_cuts::GetPID(track);
    }

    //find muon = HMN
    ccPXbox->MuonCandidate = 0x0;
    for(UInt_t itrk = 0; itrk < ntrk; itrk++){
      if(!kused[itrk]){
        if(charges[itrk]==0){
          ccPXbox->MuonCandidate = ccPXbox->TrackerTracks[itrk];
          kused[itrk]=kTRUE;

          break;
        }
      }
    }
    if(!ccPXbox->MuonCandidate){
      return kFALSE;
    }

    //find proton
    ccPXbox->ProtonCandidate = 0x0;
    for(UInt_t itrk = 0; itrk < ntrk; itrk++){
      if(!kused[itrk]){ 
        if(charges[itrk]==1 && pids[itrk]==std_numuccPX_cuts::kProton){
          //only get the first one, the second one can be set as X
          ccPXbox->ProtonCandidate = ccPXbox->TrackerTracks[itrk];
          kused[itrk]=kTRUE;

          break;
        }
      }
    }
    //due to PID can have no proton
    if(!ccPXbox->ProtonCandidate){
      return kFALSE;
    }

    //get X
    ccPXbox->XCandidate = 0x0;
    ccPXbox->XCharge = 0;
    for(UInt_t itrk = 0; itrk < ntrk; itrk++){
      if(!kused[itrk]){
        //get the first one, should not have two, will confirm later
        ccPXbox->XCandidate=ccPXbox->TrackerTracks[itrk];
        ccPXbox->XPID = pids[itrk];
        if(ccPXbox->XCandidate->Charge>0){
          ccPXbox->XCharge = 1;
        }
        else{
          ccPXbox->XCharge = -1;
        }
        kused[itrk]=kTRUE;     

        break;
      }
    }

    if((UInt_t)(kused[0]+kused[1]+kused[2]) != ntrk){
      printf("FindCandidatesCut::Apply bad ntrk sum bad kused %d %d %d ntrk %d\n", kused[0], kused[1], kused[2], ntrk); exit(1);
    }

    //setting ==

    //for ECal veto
    std_numuccPX_actions::FindECalActivity(event, box);

    //set trueCandidate
    std_numuccPX_actions::FindTrueCandidate(box);

    return kTRUE;
}


void std_numuccPX_actions::FindECalActivity(const AnaEventB& event, AnaBoxB& box)
{

  AnaBoxCCPX * ccPXbox = static_cast<AnaBoxCCPX*>(&box);

  AnaTrackB* ecaltracks[100];
  const Int_t nallecal =  anaUtils::GetAllTracksUsingECAL(event, ecaltracks);
    
  ccPXbox->NECalTracks = nallecal;
  for(Int_t iecal=0; iecal < nallecal; iecal++){
    const AnaTrackB* tmptrk = ecaltracks[iecal];
    if(!tmptrk){// || !cutUtils::CommonVertexCut(*(ccPXbox->MuonCandidate), *(tmptrk), CommonVertexCut::_cutX(), CommonVertexCut::_cutY(), CommonVertexCut::_cutZ())){
      ccPXbox->NECalTracks--;
      continue;
    }

    if(
       tmptrk == ccPXbox->MuonCandidate ||
       tmptrk == ccPXbox->ProtonCandidate ||
       tmptrk == ccPXbox->XCandidate
       ){
      ccPXbox->NECalTracks--;
    }
    else{
      //can possibly be overwritten, that is fine
      AnaTrueTrackB* tt = tmptrk->TrueTrack;
      ccPXbox->ECalPDGSim = tt? tt->PDG : std_numuccPX_cuts::kOther;
    }
  }
  
}

void std_numuccPX_actions::FindTrueCandidate(AnaBoxB& box)
{
  AnaBoxCCPX * ccPXbox = static_cast<AnaBoxCCPX*>(&box);

  const AnaTrueVertex *vtx = static_cast<AnaTrueVertex*>(box.Vertex->TrueVertex);

  //https://repo.nd280.org/viewvc/ND280/highland2/highlandUtils/src/CategoriesUtils.cxx?revision=1.7&view=markup#l1

  Int_t nhasparent = 0;

  Int_t nmuon = 0;
  Int_t nproton=0;
  Int_t npionplus=0;
  Int_t npionminus=0;
  Int_t nother=0;

  Int_t otherpdg = -999;

  for(Int_t ii=0; ii<vtx->nTrueTracks; ii++){
    AnaTrueTrackB* trueTrack = vtx->TrueTracks[ii];
    if(!trueTrack){
      continue;
    }

    if(trueTrack->ParentPDG !=0 ){
      nhasparent++;
      continue;
    }

    if(trueTrack->PDG== 13){
      nmuon++;
    }
    else if(trueTrack->PDG==2212){
      nproton++;
    }
    else if(trueTrack->PDG==211){
      npionplus++;
    }
    else if(trueTrack->PDG== -211){
      npionminus++;
    }
    else{
      otherpdg = trueTrack->PDG;
      nother++;
    }
  }

  if(nmuon==0){
    ccPXbox->EventTypeSim = AnaBoxCCPX::k0Muon;
  }
  else if(nmuon==1){
    if(nproton==0){
      ccPXbox->EventTypeSim = AnaBoxCCPX::k0Proton;
    }
    else if(nproton==1){
      if(nother==0){
        if(     npionplus==0 && npionminus==0){
          ccPXbox->EventTypeSim = AnaBoxCCPX::kZero;
        }
        else if(npionplus==1 && npionminus==0){
          ccPXbox->EventTypeSim = AnaBoxCCPX::kPionPlus;
        }
        else if(npionplus==0 && npionminus==1){
          ccPXbox->EventTypeSim = AnaBoxCCPX::kPionMinus;
        }
        else{
          ccPXbox->EventTypeSim = AnaBoxCCPX::kMPion;
        }
      }
      else{
        ccPXbox->EventTypeSim = AnaBoxCCPX::kOther;
        ccPXbox->SimOtherPDG = otherpdg;
      }
    }
    else if(nproton==2){
      ccPXbox->EventTypeSim = AnaBoxCCPX::kProton;
    }
    else{
      ccPXbox->EventTypeSim = AnaBoxCCPX::kMProton;
    }
  }
  else{
    ccPXbox->EventTypeSim = AnaBoxCCPX::kMMuon;
  }
}

bool CommonVertexCut::Apply(const AnaEventB& event, AnaBoxB& box) const
{
    (void)event;

    AnaBoxCCPX * ccPXbox = static_cast<AnaBoxCCPX*>(&box); 

    if (!ccPXbox->MuonCandidate || !ccPXbox->ProtonCandidate){
      printf("CommonVertexCut::Apply bad null muon or null proton, should not happen here. There is FindCandidatesCut!\n"); exit(1);
    }

    const Bool_t passMuP = cutUtils::CommonVertexCut(*(ccPXbox->MuonCandidate), *(ccPXbox->ProtonCandidate), CommonVertexCut::_cutX(), CommonVertexCut::_cutY(), CommonVertexCut::_cutZ());

    Bool_t passMuX = kTRUE;
    Bool_t passPX  = kTRUE;

    if(ccPXbox->XCandidate){
      if(!cutUtils::CommonVertexCut(*(ccPXbox->MuonCandidate),   *(ccPXbox->XCandidate), CommonVertexCut::_cutX(), CommonVertexCut::_cutY(), CommonVertexCut::_cutZ())){
        passMuX = kFALSE;
      }

      if(!cutUtils::CommonVertexCut(*(ccPXbox->ProtonCandidate), *(ccPXbox->XCandidate), CommonVertexCut::_cutX(), CommonVertexCut::_cutY(), CommonVertexCut::_cutZ())){
        passPX = kFALSE;
      }
    }

    return passMuP && passMuX && passPX;
}

//=================================================================================================================
//quasi-stable
//=================================================================================================================

Int_t std_numuccPX_cuts::GetPID(AnaTrackB * trk)
{
  if(trk->Charge<0){
    if(IsPion(trk)){
      return kPion;
    }
    else{
      return kOther;
    }
  }
  else{//charge>0
    if(IsPion(trk) && !IsProton(trk)){
      return kPion;
    }
    //pion exclusion makes 0.2% less, tested on full Neut MC 7002 files
    else if(!IsPion(trk) && IsProton(trk)){
    //else if(IsProton(trk)){//only for testing to recover CCQE, confirmed consistent r3481 D4
      return kProton;
    }
    else{
      return kOther;
    }
  }
}

Bool_t std_numuccPX_cuts::IsProton(AnaTrackB * trk)
{
  if(trk && anaUtils::GetPIDLikelihood(*(trk),2)>0.5)
    return kTRUE;
  else
    return kFALSE;
}

//copied from numuCCMultiPiUtils::TPCpionSelection
Bool_t std_numuccPX_cuts::IsPion(AnaTrackB * track)
{
  if(!track)
    return kFALSE;

  if ( anaUtils::GetPIDLikelihood(*track,3) < 0.3  ) return false;

  double cut1 = (anaUtils::GetPIDLikelihood(*track,0)+anaUtils::GetPIDLikelihood(*track,3))/(1.-anaUtils::GetPIDLikelihood(*track,2));

  if( track->Momentum < 500.  && cut1 < 0.8 ) return false;

  return true;
}



//=================================================================================================================
//very stable
//=================================================================================================================

numuCCPXSelection::numuCCPXSelection(bool forceBreak): SelectionBase(forceBreak) 
{
  // define the steps
  DefineSteps();  
}

bool numuCCPXSelection::FillEventSummary(const AnaBoxB& box, AnaEventB& event, Int_t allCutsPassed[])
{
    (void)box;

    // Should set event sample enumeration here
    if(allCutsPassed[0]) event.Summary->EventSample = gkSampleEnumCCPX;
    return (event.Summary->EventSample != nd280Samples::kUnassigned);
}

bool MichelElectronCut::Apply(const AnaEventB& event, AnaBoxB& box) const
{
    (void)event;

    // Temporarily, a Float_t cut to make sure it works with OldFlatTree, FlatTree and oaAnalysis
    if (box.nFGD1MichelElectrons == 0)
        return true;
    else
        return false;
}

bool FindMichelElectronsAction::Apply(const AnaEventB& event, AnaBoxB& box) const
{
    cutUtils::FindFGDMichelElectrons(event,box);

    return true;
}

bool FindTrackerTracksAction::Apply(const AnaEventB& event, AnaBoxB& box) const
{
    AnaBoxCCPX * ccPXbox = static_cast<AnaBoxCCPX*>(&box); 

    // all Tracker tracks
    ccPXbox->TrackerTracks.clear();

    //both affect size of box->TrackerTracks, which will be cut by TrackCharge
    std_numuccPX_actions::FindTracksWithTPC(event,*ccPXbox);
    //needed!! otherwise the selected sample is different!! 
    std_numuccPX_actions::FindTracksWithFGDAndNoTPC(event, *ccPXbox);

    return true;
}

void std_numuccPX_actions::FindTracksWithTPC(const AnaEventB& event, AnaBoxCCPX& ccPXbox)
{
  // all tracks using TPC
  AnaTrackB* selTracks[100];
  int nTPC = anaUtils::GetAllTracksUsingTPC(event, selTracks);

  //loop over tpc tracks
  for (Int_t i=0;i<nTPC; ++i){
    AnaTrackB* track = selTracks[i];

    // all Tracker tracks
    ccPXbox.TrackerTracks.push_back(track);
  }
}

void std_numuccPX_actions::FindTracksWithFGDAndNoTPC(const AnaEventB& event, AnaBoxCCPX& ccPXbox)
{
  // Get all tracks with FGD and no TPC
  AnaTrackB* allFgdTracks[100];
  int nFgdTracks= anaUtils::GetAllTracksUsingFGDAndNoTPC(event, allFgdTracks);

  //loop over tracks
  for (int i=0;i<nFgdTracks;i++){
    AnaTrackB* track = allFgdTracks[i];

    // All Tracker tracks
    ccPXbox.TrackerTracks.push_back(track);
  }
}

